# side_channel_demo.py
import time
import secrets

# 导入你的加密算法实现和辅助工具
from publickey.rsa_core import RSA
from publickey.elgamal_core import ElGamal
from publickey.ecc_core import get_curve_by_name, generate_keypair as ecc_gen_keys, ecc_encrypt_simplified, ecc_decrypt_simplified
from publickey.math_utils import power

# ==============================================================================
# 1. 模拟攻击的辅助函数
# ==============================================================================

def _get_key_with_specific_weight(length_bits, target_weight):
    """辅助函数: 生成一个指定位长和汉明权重(二进制中1的个数)的密钥。"""
    if target_weight > length_bits:
        raise ValueError("权重不能大于位长")
    
    bits = [1] * target_weight + [0] * (length_bits - target_weight)
    secrets.SystemRandom().shuffle(bits)
    
    if bits[0] == 0:
        for i in range(1, length_bits):
            if bits[i] == 1:
                bits[0], bits[i] = bits[i], bits[0]
                break
    return int("".join(map(str, bits)), 2)

# ==============================================================================
# 2. 针对每个算法的独立时序攻击模拟
# ==============================================================================

def simulate_rsa_timing_attack():
    """模拟针对RSA解密过程的时序攻击漏洞。"""
    print("\n" + "="*20, "模拟: RSA 解密过程的时序攻击", "="*20)
    
    key_size = 2048
    num_iterations = 100

    # 1. 创建一个正常的RSA实例，仅用于生成合法的n和e，以及加密消息
    victim_rsa = RSA(bits=key_size)
    n, e = victim_rsa.public_key
    ciphertexts = [victim_rsa.encrypt(secrets.randbits(256).to_bytes(32, 'big')) for _ in range(num_iterations)]
    print(f"已使用合法的RSA-{key_size}公钥生成 {num_iterations} 条密文。")

    # 2. 构造两个汉明权重不同的“私钥d”用于计时
    key_low_weight = _get_key_with_specific_weight(key_size, key_size // 4)
    key_high_weight = _get_key_with_specific_weight(key_size, key_size * 3 // 4)
    print("已生成高、低两种汉明权重的私钥d用于模拟解密。")

    # 3. 计时使用低权重私钥的解密操作
    start_time = time.perf_counter()
    for c in ciphertexts:
        # 我们只关心核心的模幂运算，这是解密中最耗时的部分
        power(c, key_low_weight, n)
    time_low = time.perf_counter() - start_time
    print(f"使用低权重私钥 'd' 解密 {num_iterations} 条消息的总时间: {time_low:.6f} 秒")

    # 4. 计时使用高权重私钥的解密操作
    start_time = time.perf_counter()
    for c in ciphertexts:
        power(c, key_high_weight, n)
    time_high = time.perf_counter() - start_time
    print(f"使用高权重私钥 'd' 解密 {num_iterations} 条消息的总时间: {time_high:.6f} 秒")
    
    if time_high > time_low:
        diff = (time_high - time_low) / time_low * 100
        print(f"++ 成功观察到时序差异: RSA解密时间在高权重私钥下比低权重长 {diff:.2f}%。")
    else:
        print("!! 未观察到预期的时序差异（可能由系统噪音导致）。")

def simulate_elgamal_timing_attack():
    """模拟针对ElGamal解密过程的时序攻击漏洞。"""
    print("\n" + "="*20, "模拟: ElGamal 解密过程的时序攻击", "="*20)

    key_size = 2048
    num_iterations = 100

    # 1. 创建一个正常的ElGamal实例，用于生成参数和加密
    victim_elgamal = ElGamal(key_size_bits=key_size)
    pub_key = victim_elgamal.public_key
    p = pub_key[0]
    messages_int = [secrets.randbelow(p) for _ in range(num_iterations)]
    ciphertexts = [ElGamal.encrypt(pub_key, m) for m in messages_int]
    print(f"已使用合法的ElGamal-{key_size}公钥生成 {num_iterations} 条密文。")

    # 2. 构造两个汉明权重不同的“私钥x”
    key_low_weight = _get_key_with_specific_weight(key_size, key_size // 4)
    key_high_weight = _get_key_with_specific_weight(key_size, key_size * 3 // 4)
    print("已生成高、低两种汉明权重的私钥x用于模拟解密。")

    # 3. 计时使用低权重私钥的解密操作
    start_time = time.perf_counter()
    for c1, c2 in ciphertexts:
        # 解密的核心耗时操作
        power(c1, key_low_weight, p)
    time_low = time.perf_counter() - start_time
    print(f"使用低权重私钥 'x' 解密 {num_iterations} 条消息的总时间: {time_low:.6f} 秒")

    # 4. 计时使用高权重私钥的解密操作
    start_time = time.perf_counter()
    for c1, c2 in ciphertexts:
        power(c1, key_high_weight, p)
    time_high = time.perf_counter() - start_time
    print(f"使用高权重私钥 'x' 解密 {num_iterations} 条消息的总时间: {time_high:.6f} 秒")
    
    if time_high > time_low:
        diff = (time_high - time_low) / time_low * 100
        print(f"++ 成功观察到时序差异: ElGamal解密时间在高权重私钥下比低权重长 {diff:.2f}%。")
    else:
        print("!! 未观察到预期的时序差异（可能由系统噪音导致）。")


def simulate_ecc_timing_attack():
    """模拟针对ECC解密过程的时序攻击漏洞。"""
    print("\n" + "="*20, "模拟: ECC 解密过程的时序攻击", "="*20)
    
    curve_name = "secp256r1"
    num_iterations = 100
    curve = get_curve_by_name(curve_name)
    key_size = curve.n.bit_length()

    # 1. 创建一个正常的密钥对，用于加密
    _, pub_key = ecc_gen_keys(curve)
    messages = [secrets.token_bytes(32) for _ in range(num_iterations)]
    # 加密得到 R 点列表，这是攻击的目标
    ciphertexts_R = [ecc_encrypt_simplified(pub_key, curve, m)[0] for m in messages]
    print(f"已使用合法的ECC-{curve_name}公钥生成 {num_iterations} 条密文。")
    
    # 2. 构造两个汉明权重不同的“私钥d”
    key_low_weight = _get_key_with_specific_weight(key_size, key_size // 4)
    key_high_weight = _get_key_with_specific_weight(key_size, key_size * 3 // 4)
    print("已生成高、低两种汉明权重的私钥d用于模拟解密。")
    
    # 3. 计时使用低权重私钥的解密操作
    start_time = time.perf_counter()
    for R in ciphertexts_R:
        # 解密的核心耗时操作是标量乘法
        curve.scalar_multiply(key_low_weight, R)
    time_low = time.perf_counter() - start_time
    print(f"使用低权重私钥 'd' 解密 {num_iterations} 条消息的总时间: {time_low:.6f} 秒")

    # 4. 计时使用高权重私钥的解密操作
    start_time = time.perf_counter()
    for R in ciphertexts_R:
        curve.scalar_multiply(key_high_weight, R)
    time_high = time.perf_counter() - start_time
    print(f"使用高权重私钥 'd' 解密 {num_iterations} 条消息的总时间: {time_high:.6f} 秒")

    if time_high > time_low:
        diff = (time_high - time_low) / time_low * 100
        print(f"++ 成功观察到时序差异: ECC解密时间在高权重私钥下比低权重长 {diff:.2f}%。")
    else:
        print("!! 未观察到预期的时序差异（可能由系统噪音导致）。")


if __name__ == '__main__':
    print("#"*20, " 边信道攻击：时序漏洞分算法模拟 ", "#"*20)
    print("本系列实验旨在证明，由于底层数学函数的实现非恒定时间，")
    print("RSA, ElGamal, 和 ECC 的解密操作均存在可被观察到的时序差异。")
    
    simulate_rsa_timing_attack()
    simulate_elgamal_timing_attack()
    simulate_ecc_timing_attack()